import cv2
from .color import get_colors
from .tracking import track_objects
from .frame_processor import process_frame as process_frame_base
from deep_sort_realtime.deepsort_tracker import DeepSort

def process_frame(frame, seg_model, emotion_model):
    """处理单个视频帧并返回处理后的帧和检测到的情绪"""
    try:
        # 初始化追踪器
        tracker = DeepSort(max_age=5, n_init=2)
        colors = get_colors(10)
        
        # 处理帧
        detections = process_frame_base(frame, seg_model, emotion_model)
        
        # 处理图像，添加标注
        processed_frame = track_objects(frame, detections, colors, tracker)
        
        # 获取情绪信息
        tracks = tracker.update_tracks(detections, frame=frame)
        emotions = []
        
        for track in tracks:
            if track.is_confirmed():
                emotion_label = track.det_class
                # 将情绪标签转换为中文
                emotion_chinese = {
                    'Happy': '高兴',
                    'Sad': '伤心',
                    'Angry': '生气',
                    'Surprise': '惊讶',
                    'Neutral': '中性',
                    'Calm': '平静',
                    'Fear': '恐惧',
                    'Disgust': '厌恶',
                    'Contempt': '蔑视',
                    'Unknown': '未知'
                }.get(emotion_label, '未知')
                emotions.append(emotion_chinese)
        
        return processed_frame, emotions
        
    except Exception as e:
        print(f"处理视频帧时出错: {str(e)}")
        return frame, []

def process_video(video_capture, seg_model, emotion_model):
    """处理整个视频文件"""
    all_results = []
    all_id_emotions = []
    
    tracker = DeepSort(max_age=5, n_init=2)
    colors = get_colors(10)
    
    while True:
        ret, frame = video_capture.read()
        if not ret:
            break
            
        try:
            # 处理当前帧
            processed_frame, emotions = process_frame(frame, seg_model, emotion_model)
            all_results.append(processed_frame)
            all_id_emotions.append(emotions)
            
        except Exception as e:
            print(f"处理视频帧时出错: {str(e)}")
            all_results.append(frame)
            all_id_emotions.append({})
    
    return all_results, all_id_emotions